from sklearn.model_selection import train_test_split
from sklearn.datasets import make_classification
import numpy as np
import math
import torch.nn as nn
import torch.nn.functional as F
import torch


def add_noise(y_data, noise_config: tuple = None):
    if noise_config:
        assert len(noise_config) == 2, "add_noise must be a tuple of (mean, std)"
        (mean, std) = noise_config
        noise = torch.normal(mean, std, size=y_data.shape)
        y_data += noise
    return y_data


def generate_logic_data(size: (int, tuple) = (1000, 2), noise_config: tuple = None):
    if isinstance(size, int):
        size = (size, 2)
    assert len(size) == 2, "size must be int or a tuple of length 2"
    x_data = torch.randn(size)
    y_data = torch.logical_and(x_data[:, 0] > 0, x_data[:, 1] > 0).float()

    y_data = add_noise(y_data, noise_config)

    return x_data, y_data


def generate_add_data(size: (int, tuple) = (1000, 2), noise_config: tuple = None, threshold: float = 0.0):
    if isinstance(size, int):
        size = (size, 2)
    assert len(size) == 2, "size must be int or a tuple of length 2"
    x_data = torch.randn(size)
    y_data = (x_data[:, 0] + x_data[:, 1] > threshold).float()

    y_data = add_noise(y_data, noise_config)

    return x_data, y_data


def conv_tensor_to_numpy(data):
    if isinstance(data, torch.Tensor):
        data = data.numpy()
    return data


def split_train_and_test(x_data, y_data, test_size=0.2, random_state=None) -> torch.Tensor:
    x_data = conv_tensor_to_numpy(x_data)
    y_data = conv_tensor_to_numpy(y_data)
    res = train_test_split(x_data, y_data, test_size=test_size, random_state=random_state)
    for i in range(len(res)):
        res[i] = torch.from_numpy(res[i])
    return res


if __name__ == '__main__':
    # x_data, y_data = generate_logic_data(1000)
    x_data, y_data = generate_add_data(1000)
    x_train, x_test, y_train, y_test = split_train_and_test(x_data, y_data, test_size=0.5, random_state=1)
    print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)


